package cn.com.basic.framework.datasource.aop;

import cn.com.basic.framework.datasource.annotations.DynamicDataSourceSwitch;
import cn.com.basic.framework.datasource.config.DynamicDataSourceContextHolder;
import cn.com.basic.framework.datasource.enums.DataSourceType;
import cn.com.basic.framework.datasource.util.AnnotationMetaDataBeanUtil;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.After;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.boot.context.event.ApplicationReadyEvent;
import org.springframework.context.event.EventListener;
import org.springframework.core.annotation.AnnotationUtils;
import org.springframework.stereotype.Component;
import org.springframework.util.ClassUtils;
import org.springframework.util.ReflectionUtils;

import java.lang.reflect.Method;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

@Aspect
@Component
@Slf4j
public class DynamicDataSourceHandlerAsepect {
    private static final ConcurrentHashMap<String, DataSourceType> cacheMap =
            new ConcurrentHashMap<>();

    @Pointcut("execution(* cn.com.basic.framework..*(..))")
    public void pointCut() {

    }

    /**
     * 这里可以保证所有类都以被窗口加载完成
     */
    @EventListener(ApplicationReadyEvent.class)
    public void init() {
        Set<String> annotationClass = AnnotationMetaDataBeanUtil
                .findPackageAnnotationClass("com.example.service.impl", "com.example.mapper");
        annotationClass.parallelStream().map(className -> ClassUtils.resolveClassName(className, null))
                .forEach(clazz -> ReflectionUtils.doWithLocalMethods(clazz, method -> {
                    DynamicDataSourceSwitch annotation =
                            Optional.ofNullable(AnnotationUtils.findAnnotation(method, DynamicDataSourceSwitch.class))
                                    .orElseGet(() -> AnnotationUtils.findAnnotation(clazz, DynamicDataSourceSwitch.class));
                    Optional.ofNullable(annotation).map(DynamicDataSourceSwitch::value)
                            .ifPresent(dataSourceType -> cacheMap.put(fullMethodName(clazz, method), dataSourceType));
                }));
        System.out.println(cacheMap);
    }

    public String fullMethodName(Class<?> clazz, Method method) {
        return clazz.getName() + "#" + method.getName();
    }

    @Before("pointCut()")
    public void doBefore(JoinPoint joinPoint) {
        Method method = ((MethodSignature) joinPoint.getSignature()).getMethod();
        Class<?> clazz = joinPoint.getTarget().getClass().getName().startsWith("com.sun.proxy") ?
                method.getDeclaringClass() : joinPoint.getTarget().getClass();
        // 当获取不到时，向Map中加入DefaultTargetDataSource
        DataSourceType dataSourceType = cacheMap.computeIfAbsent(fullMethodName(clazz, method),
                methodName -> Optional.ofNullable(
                        Optional.ofNullable(AnnotationUtils.findAnnotation(method, DynamicDataSourceSwitch.class))
                                .orElseGet(() -> AnnotationUtils.findAnnotation(clazz, DynamicDataSourceSwitch.class))
                ).map(DynamicDataSourceSwitch::value).orElse(DataSourceType.MASTER));
        Optional.ofNullable(dataSourceType).ifPresent(DynamicDataSourceContextHolder::setDataSource);
    }

    /**
     * 理掉当前设置的数据源，让默认的数据源不受影响
     */
    @After("pointCut()")
    public void after() {
        DynamicDataSourceContextHolder.clearDataSource();
    }

}
